{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "508244a3-05a3-4e97-8867-4c5cf3956118",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "# LiteLLM Tool calling capabilities with Unity Catalog\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "**API Key**\n",
    "To run this tutorial, you will need an OpenAI API key. \n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Below, we validate that this key is set properly in your environment.\n",
    "\n",
    "**Packages**\n",
    "\n",
    "To interface with both UnityCatalog and LiteLLM, you will need to install the following packages:\n",
    "\n",
    "```\n",
    "pip install litellm unitycatalog-client unitycatalog-litellm -q -U\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9bdec3f7-edf3-42f5-a742-e8f16c3b53fe",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "assert \"OPENAI_API_KEY\" in os.environ, (\n",
    "    \"Please set the OPENAI_API_KEY environment variable to your OpenAI API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "628b3307-4adb-4635-ad7d-c8209ec2f8ee",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Configuration and Client setup\n",
    "\n",
    "In order to connect to your Unity Catalog server, you'll need an instance of the `ApiClient` from the `unitycatalog-client` package. \n",
    "\n",
    "> Note: If you don't already have a Catalog and a Schema created, be sure to create them before running this notebook and adjust the `CATALOG` and `SCHEMA` variables below to suit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "127720ad-c167-4e58-b74b-294501b5b0a7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from unitycatalog.ai.core.client import UnitycatalogFunctionClient\n",
    "from unitycatalog.client import ApiClient, Configuration\n",
    "\n",
    "config = Configuration(host=\"http://localhost:8080/api/2.1/unity-catalog\")\n",
    "\n",
    "api_client = ApiClient(configuration=config)\n",
    "\n",
    "client = UnitycatalogFunctionClient(api_client=api_client)\n",
    "\n",
    "# Replace with your own catalog and schema\n",
    "CATALOG = \"AICatalog\"\n",
    "SCHEMA = \"AISchema\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "55a67715-63f8-4704-a284-2df3c580cf7e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Define functions and register them to Unity Catalog\n",
    "\n",
    "In this next section, we'll be defining a Python function and creating it within Unity Catalog so that it can be retrieved and used as tools within our LiteLLM call. \n",
    "\n",
    "There are a few things to keep in mind when creating functions for use with the `create_python_function` API:\n",
    "\n",
    "- Ensure that your have properly defined types for all arguments and for the return of the function.\n",
    "- Ensure that you have a Google-style docstring defined that includes descriptions for the function, each argument, and the return of the function. This is critical, as these are used to populate the metadata associated with the function within Unity Catalog, providing contextual data for an LLM to understand when and how to call the tool associated with this function.\n",
    "- If there are packages being called that are not part of core Python, ensure that the import statements are locally scoped (defined within the function body)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f0b16e6b-3c36-4bbf-bc4b-50fc81fe2bdb",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Define our function\n",
    "def get_weather(location: str) -> str:\n",
    "    \"\"\"\n",
    "    Query the weather and return a string with the current temperature.\n",
    "\n",
    "    Args:\n",
    "      location: The location to get the weather\n",
    "\n",
    "    \"\"\"\n",
    "    return \"Hot!\"\n",
    "\n",
    "\n",
    "# Register functions\n",
    "function_info = client.create_python_function(\n",
    "    func=get_weather, catalog=CATALOG, schema=SCHEMA, replace=True\n",
    ")\n",
    "execution_function_name = function_info.full_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "24509e03-089e-4673-882d-9d6d1152555d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Create a Toolkit instance of the functions\n",
    "\n",
    "Now that the functions have been created within Unity Catalog, we can use the `unitycatalog-litellm` package to create a toolkit instance that our Agent will 'understand' as valid tools to use within its APIs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2675f057-e152-45fe-8a38-b710cf23b3d4",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[LiteLLMTool(name='main__default__get_weather', description='Query the weather and return a string with the current temperature.', tool={'type': 'function', 'function': {'name': 'main__default__get_weather', 'strict': True, 'parameters': {'properties': {'location': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'description': 'The location to get the weather', 'title': 'Location'}}, 'title': 'main__default__get_weather__params', 'type': 'object', 'additionalProperties': False, 'required': ['location']}, 'description': 'Query the weather and return a string with the current temperature.'}})]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from unitycatalog.ai.litellm.toolkit import UCFunctionToolkit\n",
    "\n",
    "toolkit = UCFunctionToolkit(function_names=[execution_function_name])\n",
    "tools = toolkit.tools\n",
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b4cd8af9-9b53-4c13-97f2-75550d5a45e6",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Pass a Tool Definition to LiteLLM\n",
    "\n",
    "Now that we have a UCFunctionToolkit with definitions of our tool, we can leverage them in LiteLLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "102272ba-fd86-414a-9bed-d08502c2ff50",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ModelResponse(id='chatcmpl-Ae2ploBk7zcxPh16PHWaJXJiI9qRa', created=1734107377, model='gpt-4o-mini-2024-07-18', object='chat.completion', system_fingerprint='fp_6fc10e10eb', choices=[Choices(finish_reason='tool_calls', index=0, message=Message(content=None, role='assistant', tool_calls=[ChatCompletionMessageToolCall(function=Function(arguments='{\"location\":\"San Francisco\"}', name='main__default__get_weather'), id='call_KVNXgt8RRqUjug7iaTnOQD7d', type='function')], function_call=None))], usage=Usage(completion_tokens=19, prompt_tokens=84, total_tokens=103, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier=None)\n"
     ]
    }
   ],
   "source": [
    "import litellm\n",
    "\n",
    "# Define your request\n",
    "question = \"What is the current temperature in SF?\"\n",
    "messages = [{\"role\": \"user\", \"content\": question}]\n",
    "\n",
    "# Show the response\n",
    "response = litellm.completion(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=messages,\n",
    "    tools=tools,\n",
    "    tool_choice=\"auto\",  # auto is default, but we'll be explicit\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "44ba48bb-0986-49ec-be5e-71e0465ed207",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Call the UC Function\n",
    "Based on the above response, LiteLLM wants to call the tool. Let's execute the function via the `generate_tool_call_messages` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "569a79fa-5ad5-44b0-a98e-0837c45c7e12",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " ModelResponse(id='chatcmpl-Ae2povzf9i8bRpqwCRXBq3YsbP2kH', created=1734107380, model='gpt-4o-mini-2024-07-18', object='chat.completion', system_fingerprint='fp_6fc10e10eb', choices=[Choices(finish_reason='stop', index=0, message=Message(content='The current temperature in San Francisco is hot! For the exact temperature, you might want to check a reliable weather website or app.', role='assistant', tool_calls=None, function_call=None))], usage=Usage(completion_tokens=26, prompt_tokens=48, total_tokens=74, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier=None)\n"
     ]
    }
   ],
   "source": [
    "from unitycatalog.ai.litellm.utils import generate_tool_call_messages\n",
    "\n",
    "tool_messages = generate_tool_call_messages(\n",
    "    response=response, client=client, conversation_history=messages\n",
    ")\n",
    "\n",
    "response_2 = litellm.completion(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=tool_messages,\n",
    ")\n",
    "response_2"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "computePreferences": null,
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "client": "1"
   },
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 2
   },
   "notebookName": "UC LiteLLM OSS (on databricks)",
   "widgets": {}
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
